package fem2.strategies;

import fem2.Debugger;
import fem2.Model;
import fem2.enu.ConvergenceCriteria;
import fem2.enu.EchoLevelType;
import fem2.enu.LineSearchCriteria;
import fem2.jlinalg.Solver;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IVector;
import inf.jlinalg.MatrixFormat;
import inf.jlinalg.SolveFailedException;
import inf.math.LinAlg;

import java.util.Arrays;

import math2.MathUtilities;

/**
 * Implementation of Newton-Raphson solution strategy with built in line search
 * algorithm. Default line search criteria is NO_LINE_SEARCH
 * 
 * @author hbui
 * 
 */
public abstract class NewtonRaphsonStrategy extends Strategy {

	protected ArrayVector r = new ArrayVector();
	private double nr;
	protected ArrayVector uHat = null;
	protected LineSearchCriteria lineSearchCriteria;

	public NewtonRaphsonStrategy(Model m) {
		super(m);
		setLineSearchCriteria(LineSearchCriteria.NO_LINE_SEARCH);
	}

	public NewtonRaphsonStrategy(Model m, LineSearchCriteria criteria) {
		super(m);
		setLineSearchCriteria(criteria);
	}

	/**
	 * set line search criteria for this solution strategy
	 * 
	 * @param lineSearchCriteria
	 */
	public void setLineSearchCriteria(LineSearchCriteria lineSearchCriteria) {
		this.lineSearchCriteria = lineSearchCriteria;
	}

	/**
	 * 
	 * @return the next lambda for time step iteration
	 */
	public abstract double getNextTimeStep();

	@Override
	public double[] getCurrentUHat() {
		return uHat.getValues();
	}

	@Override
	public double[] getCurrentUdHat() {
		throw new Error("this method is not available for this strategy");
	}

	@Override
	public double[] getCurrentUddHat() {
		throw new Error("this method is not available for this strategy");
	}

	@Override
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * initialize model (enumeration)
		 */
		super.started(leftHandSideSolver, rightHandSide, u);
		int n = model.getSize();

		/*
		 * initialize solver
		 */
		initializeSolver(leftHandSideSolver);

		/*
		 * initialize auxiliary vector
		 */
		r.setSize(n);
		u.setSize(n);
		rightHandSide.setSize(n);
		uHat = new ArrayVector(n);
	}

	@Override
	public void timeStepStarted(Solver leftHandSide, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * increment the time step
		 */
		double t = getNextTimeStep();

		if (getEchoLevel().contain(EchoLevelType.OUTPUT)) {
			Debugger.watch("timestep start at t = " + t);
		}

		/*
		 * zero out and assemble external force vector
		 */
		Arrays.fill(r.getValues(), 0.0);
		model.assembleR(r.getValues(), t);

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("system external force vector:\n", r);
		}

		nr = BLAM.norm(BLAM.NORM_TWO, r);
		if (nr < MathUtilities.tol12) {
			if (convergenceCriteria == ConvergenceCriteria.RELATIVE_RESIDUAL) {
				convergenceCriteria = ConvergenceCriteria.ABSOLUTE_RESIDUAL;
				Debugger.warn("norm of initial external force is so small, switch to "
						+ convergenceCriteria.toString());
			}
		}

		/*
		 * set the initial solution vector and initialize model at time step
		 */
		model.setX(t, u.getValues());
		model.InitializeSolutionStep();

		/*
		 * assemble internal force vector
		 */
		model.assembleRi(rightHandSide.getValues());

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("initial system internal force vector:\n", rightHandSide);
		}

		/*
		 * compute residual
		 */
		BLAM.add(-1.0, r, rightHandSide);
	}

	@Override
	public void iterationStarted(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		model.InitalizeNonlinearIteration();

		/*
		 * zero out and assemble tangent stiffness
		 */
		leftHandSideSolver.initialize();
		model.assembleKt(leftHandSideSolver.getA());

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("system tangent stiffness:\n", leftHandSideSolver.getA());
		}

		double t = getCurrentTimeStep();

		/*
		 * solve
		 */
		try {
			leftHandSideSolver.solve(rightHandSide.getValues());
			/*
			 * at this point rightHandSide is -du
			 */
		} catch (SolveFailedException e) {
			throw new Error("solve failed at lambda = " + t + ", error: " + e.getMessage());
		}
	}

	@Override
	public boolean iterationFinished(Solver leftHandSideSolver, ArrayVector rightHandSide,
			ArrayVector u) {
		double t = getCurrentTimeStep();
		double error = 0.0;

		/*
		 * compute line search increment
		 */
		double eta = computeLineSearchStep(rightHandSide, u);

		/*
		 * update solution
		 */
		BLAM.add(-eta, rightHandSide, u);

		if (convergenceCriteria == ConvergenceCriteria.ABSOLUTE_DISPLACEMENT
				|| convergenceCriteria == ConvergenceCriteria.RELATIVE_DISPLACEMENT) {
			error = BLAM.norm(BLAM.NORM_TWO, rightHandSide);
		}

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			System.out.println("u_new = " + MatrixFormat.format(u));
		}

		/*
		 * update model
		 */
		model.setX(t, u.getValues());
		model.FinalizeNonlinearIteration(); // update material

		/*
		 * recompute internal force vector
		 */
		Arrays.fill(rightHandSide.getValues(), 0);
		model.assembleRi(rightHandSide.getValues());

		if (getEchoLevel().contain(EchoLevelType.DEBUG)) {
			System.out.println("system internal force vector: "
					+ MatrixFormat.format(rightHandSide));
		}

		if (convergenceCriteria == ConvergenceCriteria.ABSOLUTE_RESIDUAL
				|| convergenceCriteria == ConvergenceCriteria.RELATIVE_RESIDUAL) {
			/*
			 * compute residual
			 */
			BLAM.add(-1.0, r, rightHandSide);
		}

		/*
		 * check convergence criteria
		 */
		switch (convergenceCriteria) {
		case ABSOLUTE_RESIDUAL:
			error = BLAM.norm(BLAM.NORM_TWO, rightHandSide);
			break;
		case RELATIVE_RESIDUAL:
			error = BLAM.norm(BLAM.NORM_TWO, rightHandSide) / nr;
			break;
		case ABSOLUTE_DISPLACEMENT:
			/*
			 * do nothing
			 */
			break;
		case RELATIVE_DISPLACEMENT:
			double norm_du = LinAlg.diffNorm2(uHat.getSize(), u.getValues(), uHat.getValues());
			error /= norm_du;
			break;
		default:
			throw new Error("invalid convergence criteria");
		}

		// double norm_du = BLAM.norm(BLAM.NORM_TWO, rightHandSide);
		// double norm_u = BLAM.norm(BLAM.NORM_TWO, u);
		// double conv_disp = 0.0;
		// if (norm_u < MyLinAlg.tol12) {
		// conv_disp = norm_du;
		// } else {
		// conv_disp = norm_du / norm_u;
		// }

		// double residual = BLAM.norm(BLAM.NORM_TWO, rightHandSide);
		// double error;
		// if (nr < LinAlgUtilities.tol12) {
		// error = residual;
		// } else {
		// error = residual / nr;
		// }
		// error = Math.max(error, conv_disp);

		if (getEchoLevel().contain(EchoLevelType.OUTPUT)) {
			Debugger.watch("substep " + convergenceCriteria.toString() + " error = " + error);
		}

		return (error < tol);
	}

	@Override
	public void timeStepFinished(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		model.FinalizeSolutionStep();
		/*
		 * save solution vector
		 */
		BLAM.copy(u, uHat);
	}

	/**
	 * compute line search increment
	 * 
	 * @param mdu
	 *            -du
	 * @param u
	 *            u
	 * @return eta
	 */
	private double computeLineSearchStep(IVector mdu, IVector u) {

		switch (lineSearchCriteria) {

		case NO_LINE_SEARCH:
			return 1.0;

		case INTERPOLATED_LINE_SEARCH:
			double t = getCurrentTimeStep();
			ArrayVector U = (ArrayVector) u.clone();
			ArrayVector R = new ArrayVector(model.getSize());

			/*
			 * compute s0
			 */
			model.setX(t, U.getValues());
			model.FinalizeNonlinearIteration(); // update material
			model.assembleRi(R.getValues());
			BLAM.add(-1.0, r, R);
			double s0 = BLAM.dot(mdu, R);

			/*
			 * compute first s
			 */
			double eta = 1.0;
			BLAM.add(-eta, mdu, U);
			model.setX(t, U.getValues());
			model.FinalizeNonlinearIteration(); // update material
			Arrays.fill(R.getValues(), 0.0);
			model.assembleRi(R.getValues());
			BLAM.add(-1.0, r, R);
			double s = BLAM.dot(mdu, R);

			int cnt = 0;
			while (Math.abs(s / s0) > MathUtilities.tol6 && (++cnt < MAX_NT)) {
				eta = eta * s0 / (s0 - s);

				/*
				 * update U
				 */
				BLAM.copy(u, U);
				BLAM.add(-eta, mdu, U);

				/*
				 * update s
				 */
				model.setX(t, U.getValues());
				model.FinalizeNonlinearIteration(); // update material
				Arrays.fill(R.getValues(), 0.0);
				model.assembleRi(R.getValues());
				BLAM.add(-1.0, r, R);
				s = BLAM.dot(mdu, R);

				if (getEchoLevel().contain(EchoLevelType.AUX5)) {
					Debugger.watch("line search stepping =", eta);
				}
			}

			if (cnt == MAX_NT) {
				throw new Error("line search stepping does not converge in " + MAX_NT
						+ " iterations");
			}

			if (getEchoLevel().contain(EchoLevelType.OUTPUT)) {
				Debugger.watch("line search converged after " + cnt + " iterations at eta = " + eta);
			}

			return eta;
			/*
			 * TODO
			 */

		default:
			throw new Error("invalid line search criteria");
		}
	}
}
